package zhengqc.tree;
/* 平衡二叉树: 对于任意一个节点, 左右子树的高度差不能超过1 */
/* 平衡因子: 某节点的左右子树的高度差 */

import zhengqc.list.ArrayList;

public class AVLTree<E extends Comparable<E>> extends AbstractBSTree<E> {
    public AVLTree() {
        this.root = null;
        this.size = 0;
        this.tempList = new ArrayList<>();
        tempElement = null;
    }
    /* 对节点y进行向右旋转操作，返回旋转后新的根节点x
            y                              x
           / \                           /   \
          x   T4     向右旋转 (y)        z     y
         / \       - - - - - - - ->    / \   / \
        z   T3                       T1  T2 T3 T4
       / \
     T1   T2                                    */
    private Node rightRotate(Node y) {
        Node x = y.left;
        y.left = x.right;
        x.right = y;
        // 更新 height
        y.height = Math.max(getHeight(y.left), getHeight(y.right)) + 1;
        x.height = Math.max(getHeight(x.left), getHeight(x.right)) + 1;
        return x;
    }
    /* 对节点y进行向左旋转操作，返回旋转后新的根节点x
        y                             x
      /  \                          /   \
     T1   x      向左旋转 (y)       y     z
         / \   - - - - - - - ->   / \   / \
       T2  z                     T1 T2 T3 T4
          / \
         T3 T4                                  */
    private Node leftRotate(Node y) {
        Node x = y.right;
        y.right = x.left;
        x.left = y;
        // 更新 height
        y.height = Math.max(getHeight(y.left), getHeight(y.right)) + 1;
        x.height = Math.max(getHeight(x.left), getHeight(x.right)) + 1;
        return x;
    }
    private Node updateNode(Node node) {
        // 对于删除, 须有此条件判断; 对于添加, 则可有可无
        if (node == null) {
            return null;
        }
        // 更新 node 的 height
        node.height = 1 + Math.max(getHeight(node.left), getHeight(node.right));
        // 计算平衡因子
        int balanceFactor = getBalanceFactor(node);

        // 加入节点后, 沿着节点向上维护平衡性
        // 对于下面 "=" 条件, 对于删除来说是必须的, 而对于添加则是可有可无的
        // LL : 单次右旋
        if ( (balanceFactor > 1) &&
                (getBalanceFactor(node.left) >= 0) ) {
            return rightRotate(node);
        }
        // RR : 单次左旋
        if ( (balanceFactor < -1) &&
                (getBalanceFactor(node.right) <= 0) ) {
            return leftRotate(node);
        }
        // LR -> LL : 先左旋, 再右旋
        if ( (balanceFactor > 1) &&
                (getBalanceFactor(node.left) < 0) ) {
            node.left = leftRotate(node.left);
            return rightRotate(node);
        }
        // RL -> RR : 先右旋, 再左旋
        if ( (balanceFactor < -1) &&
                (getBalanceFactor(node.right) > 0) ) {
            node.right = rightRotate(node.right);
            return leftRotate(node);
        }
        return node;
    }

    /* 向以 node 为根节点的 AVLTree 树中插入元素 element, 返回插入新节点后的 AVLTree 的根节点 */
    private Node addNode(Node node, E element) {
        if (node == null) {
            ++ this.size;
            return new Node(element);
        }
        if (element.compareTo(node.element) < 0) {
            node.left = addNode(node.left, element);
        } else if (element.compareTo(node.element) > 0) {
            node.right = addNode(node.right, element);
        } else { // element.compareTo(node.element) == 0
            node.element = element; // 相当于不做操作
        }
        return updateNode(node);
    }
    @Override /* 向 AVLTree 插入元素 */
    public void add(E element) {
        if (element == null) {
            throw new IllegalArgumentException("element is null!");
        }
        this.root = addNode(this.root, element);
    }

    /* 删除以 node 为根的 AVLTree 中值为 element 的节点; 返回删除节点后新的 AVLTree 的根节点 */
    private Node removeNode(Node node, E element) {
        // 未找到值为 element 的节点
        if (node == null) {
            return null;
        }

        Node retNode = null;
        if (element.compareTo(node.element) < 0) {
            node.left = removeNode(node.left, element);
            retNode = node;
        } else if (element.compareTo(node.element) > 0) {
            node.right = removeNode(node.right, element);
            retNode = node;
        } else {
            // 先保存带删除的节点值
            this.tempElement = node.element;

            if (node.left == null) { // 待删除节点左子树为空
                -- this.size;
                retNode = node.right;
            } else if (node.right == null) { // 待删除节点右子树为空
                -- this.size;
                retNode = node.left;
            } else {
                // 待删除节点的左右子树均不为空
                // 找到比待删除节点大的最小节点, 即待删除节点右子树的最小节点, 用这个节点代替待删除节点的位置
                Node successor = getMinimumNode(node.right);
                successor.right = removeMinimumNode(node.right);
                successor.left = node.left;
                retNode = successor;
            }
        }

        return updateNode(retNode);
    }
    @Override /* 从 ALVTree 中删除指定元素; 原则上不需要返回值, 仅供特殊需求, 如实现 Map */
    public E remove(E element) {
        this.tempElement = null;
        this.root = removeNode(this.root, element);
        return this.tempElement;
    }

    /* 删除以 node 为根节点的 AVLTree 中的最小节点; 返回删除节点后新的 AVLTree的根节点 */
    private Node removeMaximumNode(Node node) {
        if (node.right == null) {
            -- this.size;
            this.tempElement = node.element;
            return node.left;
        }
        node.right = removeMaximumNode(node.right);
        return updateNode(node);
    }
    @Override /* 删除 AVLTree 中最大的元素, 并返回 */
    public E removeMaximum() {
        this.tempElement = null;
        if (this.isEmpty()) {
            throw new IllegalArgumentException("BST is empty");
        }
        this.root = removeMaximumNode(this.root);
        return this.tempElement;
    }

    /* 删除以 node 为根节点的 AVLTree 中的最小节点; 返回删除节点后新的 AVLTree 的根节点 */
    private Node removeMinimumNode(Node node) {
        if (node.left == null) {
            -- this.size;
            this.tempElement = node.element;
            return node.right;
        }
        node.left = removeMinimumNode(node.left);
        return updateNode(node);
    }
    @Override /* 删除 AVLTree 中最小的元素, 并返回 */
    public E removeMinimum() {
        this.tempElement = null;
        if (this.isEmpty()) {
            throw new IllegalArgumentException("BST is empty");
        }
        this.root = removeMinimumNode(this.root);
        return this.tempElement;
    }
}
